import asyncio
import datetime
import time

from langchain_community.callbacks import get_openai_callback
from langchain_core.callbacks import CallbackManager
from loguru import logger

from ai_engine.car_wrap.chain.chat_completion_chain import ChatCompletionChain
from ai_engine.core.base.base_service import BaseService
from ai_engine.core.callbacks.chain_stream_handler import QAChainStreamHandler
from ai_engine.core.model import result
from ai_engine.core.model.chat import CompletionRequest, ChatMessage
from ai_engine.core.vectorstore.vector_data.knowledge.knowledge_vs import KnowledgeVs

from ai_engine.utils.param_util import ParamUtil
from ai_engine.utils.validate_util import is_mobile, is_wechat_account


class ChatService(BaseService):
    """AI聊天业务处理类"""

    def _completion_validata(self, request: CompletionRequest):
        if request.model_style not in ["stable", "lively"]:
            raise ValueError(f"model_style error,Only supports stable or lively.")
        if len(request.model_name) < 1:
            raise ValueError(f"model_name error,Need to provide.")
        if len(request.prompt) < 1:
            raise ValueError(f"prompt error,Need to provide.")
        if len(request.chat_history) > 0:
            for message in request.chat_history:
                if message.role not in ["human", "assistant"]:
                    raise ValueError(f"chat_history role error,Only supports human or assistant.")
        if len(request.vector_categories) < 1:
            raise ValueError(f"vector_categories error,Need to provide.")
        for categories in request.vector_categories:
            if len(categories) < 1:
                raise ValueError(f"vector_categories error,Need to provide.")

    def completions(self, request: CompletionRequest):
        # 验证输入参数
        start_time = time.time()
        self._completion_validata(request)
        # 创建监听器
        chainStreamHandler = None
        callbackManager = None
        if request.stream:
            chainStreamHandler = QAChainStreamHandler(method_name="文本生成", request_id=request.request_id)
            callbackManager = CallbackManager([chainStreamHandler])

        # 创建回答问题语言模型
        llm = self.get_llm(
            temperature=request.temperature,
            stream=request.stream,
            verbose=True,
            max_tokens=request.max_tokens,
            callback_manager=callbackManager
        )
        knowledge_vs = KnowledgeVs(model_kwargs=self.model_kwargs)

        chain = ChatCompletionChain(llm=llm,
                                    knowledge_vs=knowledge_vs,
                                    request=request)

        prompt_parameter = self._params_to_prompt(request)

        if request.stream:
            asyncio.set_event_loop(asyncio.new_event_loop())
            asyncio.get_event_loop().run_in_executor(None, chain.invoke, prompt_parameter)
            generate = chainStreamHandler.generate_tokens()
            return generate
        else:
            with get_openai_callback() as cb:
                response = chain.invoke(prompt_parameter)
                answer = response["answer"]
                res = result.success(answer)
                end_time = time.time()
                logger.info("文本生成同步运行时间：%.2f秒" % (end_time - start_time))
                logger.info("文本生成同步应答内容：" + answer)
                logger.info("文本生成同步token消耗：" + str(cb))
                return res

    def _params_to_prompt(self, request: CompletionRequest):

        # 参数清洗
        request.role.name = request.role.name.replace("\n", "").strip()
        request.role.age = request.role.age.replace("\n", "").strip()
        request.role.sex = request.role.sex.replace("\n", "").strip()
        request.role.company = request.role.company.replace("\n", "").strip()
        request.role.company_abbreviation = request.role.company_abbreviation.replace("\n", "").strip()

        request.role.brand = request.role.brand.replace("\n", "").strip()
        request.role.industry = request.role.industry.replace("\n", "").strip()
        request.role.position = request.role.position.replace("\n", "").strip()
        request.role.characters = request.role.characters.replace("\n", "").strip()
        request.role.company_address = request.role.company_address.replace("\n", "").strip()
        request.role.company_tel = request.role.company_tel.replace("\n", "").strip()
        request.role.company_opening = request.role.company_opening.replace("\n", "").strip()
        request.role.company_liveTime = request.role.company_liveTime.replace("\n", "").strip()
        request.role.company_liveToStore_discount = request.role.company_liveToStore_discount.replace("\n", "").strip()
        request.role.company_testDrive_discount = request.role.company_testDrive_discount.replace("\n", "").strip()
        request.role.company_other = request.role.company_other.replace("\n", "").strip()

        request.custom.custom_vehicle_type = request.custom.custom_vehicle_type.replace("\n", "").strip()
        request.custom.custom_phone_number = request.custom.custom_phone_number.replace("\n", "").strip()
        request.custom.custom_province = request.custom.custom_province.replace("\n", "").strip()
        request.custom.custom_city = request.custom.custom_city.replace("\n", "").strip()
        request.custom.custom_area = request.custom.custom_area.replace("\n", "").strip()
        request.custom.custom_weChat_account = request.custom.custom_weChat_account.replace("\n", "").strip()

        request.role.company_province = request.role.company_province.replace('省', '')
        request.role.company_city = request.role.company_city.replace('市', '')

        if len(request.role.name) == 0:
            request.role.name = "客服小姐姐"
        if len(request.role.age) == 0:
            request.role.age = 25
        if len(request.role.sex) == 0:
            request.role.sex = "女"
        # 如果是完整的电话号码进行改写
        if is_mobile(request.prompt):
            request.prompt = "我的电话是" + request.prompt + ",尽快联系我"
        elif is_wechat_account(request.prompt):
            # 如果是微信号
            request.prompt = "我的微信是" + request.prompt + ",尽快联系我"

        now = datetime.datetime.now()
        yearStr = now.strftime('%Y年')
        monthStr = now.strftime('%m月').lstrip('0')
        dayStr = now.strftime('%d日').lstrip('0')
        weekdayStr = now.weekday() + 1
        timeStr = now.strftime('%H').lstrip('0')
        company = f"当前时间|今天是{yearStr}{monthStr}{dayStr},周(星期){weekdayStr},{timeStr}点钟。\n"
        # if len(request.role.company_vehicleSeries) > 0:
        #     company = company + "卖什么车|在售车系有“{0}，如果客户咨询的不是在售车系，你就说：{1}\n".format(
        #         request.role.company_vehicleSeries, on_sales_series)
        if len(request.role.company) > 0:
            company = company + "门店公司名|{0}。\n".format(request.role.company)
        if len(request.role.company_address) > 0:
            company = company + "门店地址|{0}。\n".format(request.role.company_address)
        else:
            company = company + "门店地址|安排销售薇上给您发定位。"

        if len(request.role.company_opening) > 0:
            company = company + "门店营业时间|{0}。\n".format(request.role.company_opening)
        if len(request.role.company_liveTime) > 0:
            company = company + "门店直播时间|{0}。\n".format(request.role.company_liveTime)

        chat_history = request.chat_history
        human_msg = ChatMessage(role="human", content=request.prompt)
        chat_history.append(human_msg)
        chat_history_str = ParamUtil.get_chat_history_str(chat_history)

        prompt = {
            "question": request.prompt,
            "chat_history": chat_history_str,
            "brand": request.role.brand,
            "company_abbreviation": request.role.company_abbreviation,
            "industry": request.role.industry,
            "position": request.role.position,
            "name": request.role.name,
            "sex": request.role.sex,
            "age": request.role.age,
            "characters": request.role.characters,
            "company": company,
            "company_vehicleSeries": request.role.company_vehicleSeries,
            "custom_phone_number": request.custom.custom_phone_number,
            "custom_weChat_account": request.custom.custom_weChat_account,
        }

        return prompt
